Skip to content

Latest commit

 

History

History

segmentation

InternViT-6B for Semantic Segmentation

This folder contains the implementation of the InternViT-6B for semantic segmentation.

Our segmentation code is developed on top of MMSegmentation v0.30.0.

🛠️ Installation

See INSTALLATION.md

📦 Data Preparation

Prepare the ADE20K dataset according to the guidelines in MMSegmentation.

📦 Model Preparation

model name type download size
InternViT-6B-224px pytorch 🤗 HF link 12 GB

Please download the above model weight and place it in the pretrained/ folder.

The directory structure is:

pretrained
└── intern_vit_6b_224px.pth

🔥 Training

Please note, this open-source code does not include DeepSpeed in MMSegmentation, so it currently only supports training for linear probing and head tuning, and does not support full-parameter training.

To train a linear classifier for InternViT-6B with 8 GPU on 1 node (total batch size 16), run:

sh dist_train.sh configs/intern_vit_6b/linear_probing/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.py 8
# or manage jobs with slurm
GPUS=8 sh slurm_train.sh <partition> <job-name> configs/intern_vit_6b/linear_probing/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.py

📊 Evaluation

type backbone head mIoU config download
few-shot (1/16) InternViT-6B Linear 46.5 config ckpt | log
few-shot (1/8) InternViT-6B Linear 50.0 config ckpt | log
few-shot (1/4) InternViT-6B Linear 53.3 config ckpt | log
few-shot (1/2) InternViT-6B Linear 55.8 config ckpt | log
few-shot (1/1) InternViT-6B Linear 57.2 config ckpt | log
linear probing InternViT-6B (frozen) Linear 47.2 config ckpt | log
head tuning InternViT-6B (frozen) UperNet 54.9 config ckpt | log
full tuning InternViT-6B UperNet 58.9 config ckpt | log

You can download checkpoints from here or from the table above. Then place them to segmentation/checkpoints/.

For example, to evaluate the InternViT-6B with a single GPU:

python test.py configs/intern_vit_6b/linear_probing/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.py checkpoints/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.pth --eval mIoU

For example, to evaluate the InternViT-6B with a single node with 8 GPUs:

sh dist_test.sh configs/intern_vit_6b/linear_probing/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.py checkpoints/linear_intern_vit_6b_504_80k_ade20k_bs16_lr4e-5_frozen.pth 8 --eval mIoU